import os
import cv2
import numpy as np
from myCVHelper import my_cv_helper as helper
from myCVHelper.my_cv_helper import logger
from PIL import Image, ImageEnhance
import traceback
import matplotlib
from matplotlib import pyplot as plt
import random

# 设置matplotlib正常显示中文和负号
matplotlib.rcParams['font.sans-serif'] = ['SimHei']  # 用黑体显示中文
matplotlib.rcParams['axes.unicode_minus'] = False  # 正常显示负号


class Configs:
    file_path = './imgs'
    file_exts = ['png', ]
    window_name = "process_"


def process_pre_tasks(im: np.ndarray, show=False) -> np.ndarray:
    if show:
        helper.Show.window_clear()
        helper.Show.resize = 4
        helper.Show.offset = (200, 100)
    if show:
        helper.Show.imshow(Configs.window_name + "原图", im, 0)
    # 转换到hsv处理
    im = cv2.cvtColor(im, cv2.COLOR_BGR2HLS)
    if show:
        helper.Show.imshow(Configs.window_name + "HSV", im, 0)
    # 中值滤波降噪
    im = cv2.medianBlur(im, 13)
    if show:
        helper.Show.imshow(Configs.window_name + "中值滤波", im, 0)
    # 消除高光
    _, mask = cv2.threshold(im, 100, 255, cv2.THRESH_BINARY)
    if show:
        helper.Show.imshow(Configs.window_name + "mask", mask, 0)
    im = cv2.illuminationChange(im, mask)
    if show:
        helper.Show.imshow(Configs.window_name + "消除高光", im, 1)

    # 增加对比度
    image = Image.fromarray(im)
    image = ImageEnhance.Contrast(image).enhance(3)
    im = np.array(image)
    if show:
        helper.Show.imshow(Configs.window_name + "增强对比度", im, 1)

    if show:
        helper.Controls.wait_exit(0)
    return im


def process(src: np.ndarray, use_track=False, refresh=True) -> np.ndarray:
    im = process_pre_tasks(src, show=False)
    if refresh:
        helper.Show.window_clear()
    # 初始化窗口配置
    helper.Show.resize = 2
    helper.Show.offset = (20, 10)
    color_circle = (0, 255, 0)
    color_rect = (0, 0, 255)

    # 试试先Canny
    # 对hsv做canny
    # channels = cv2.split(im)
    # logger.debug("channels[2].shape: %s" % str(np.array(channels[2]).shape))
    # canny = cv2.Canny(channels[2], 50, 200)
    # canny = cv2.Canny(im, 20, 230)
    canny = cv2.Canny(im, 300, 300)
    # helper.Show.imshow(Configs.window_name + "channels[2]", channels[2], 0)
    helper.Show.imshow(Configs.window_name + "Canny", canny, 0)

    def onchange_canny(window_name: str, image: np.ndarray, args: list):
        if len(args) != 2:
            logger.warning("len(args) == %s!" % str(len(args)))
            return
        image = cv2.Canny(image, cv2.getTrackbarPos(args[0].bar_name, args[0].window_name),
                          cv2.getTrackbarPos(args[1].bar_name, args[1].window_name))
        # cv2.imshow(window_name, image)
        helper.Show.imshow(window_name, image)

    if use_track:
        helper.Controls.adjust_x(Configs.window_name + "Canny调参", im, onchange_canny, [
            helper.Controls.ArgBase("arg1", 20, 555),
            helper.Controls.ArgBase("arg2", 230, 555)
        ])

    # 转灰度
    im = cv2.cvtColor(im, cv2.COLOR_HLS2BGR)
    im = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)
    helper.Show.imshow(Configs.window_name + "灰度", im, 0)

    # # 二值化
    # _, im = cv2.threshold(im, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    # helper.Show.imshow(Configs.window_name + "二值化", im, 0)
    #
    # im2 = im.copy()
    #
    # # 开操作
    # kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (31, 31))
    # im = cv2.morphologyEx(im, cv2.MORPH_OPEN, kernel)
    # helper.Show.imshow(Configs.window_name + "开操作", im, 0)
    #
    # # 描边
    # im = cv2.Canny(im, 50, 200)
    # helper.Show.imshow(Configs.window_name + "边缘检测", im, 1)

    # # 测试一下调参
    # def onchange(win_name: str, image: np.ndarray, val: int):
    #     _, image = cv2.threshold(image, val, 255, cv2.THRESH_BINARY)
    #     helper.Show.imshow(win_name, image, 4)

    # helper.Controls.adjust('测试', Configs.window_name + "调参测试", im, 0, 255, onchange)

    # 再模糊一下
    # draw = src.copy()
    im = cv2.medianBlur(im, 3)
    # draw = cv2.cvtColor(im.copy(), cv2.COLOR_GRAY2BGR)
    draw = src.copy()

    # 多个参数调整尝试
    def onchange_param1(win_name: str, image: np.ndarray, val: int):
        if val < 10:
            return
        circles1 = cv2.HoughCircles(im, cv2.HOUGH_GRADIENT, 2, im.shape[1] / 8, val, 50, minRadius=20, maxRadius=0)
        # print(circles)
        draw1 = cv2.cvtColor(image.copy(), cv2.COLOR_GRAY2BGR)
        for c in circles1[0, :]:
            try:
                if len(c) != 3: continue
            except TypeError:
                continue
            c = list(map(lambda x: int(x), c))
            # print(c)
            cv2.circle(draw1, (c[0], c[1]), 1, (0, 255, 0), 2)
            cv2.circle(draw1, (c[0], c[1]), c[2], (0, 128, 255), 2)
        helper.Show.imshow(win_name, draw1, 2)
        return draw1

    def onchange_param2(win_name: str, image: np.ndarray, val: int):
        if val < 10:
            return
        circles1 = cv2.HoughCircles(im, cv2.HOUGH_GRADIENT, 2, im.shape[1] / 8, 100, val, minRadius=20, maxRadius=0)
        # print(circles)
        draw1 = cv2.cvtColor(image.copy(), cv2.COLOR_GRAY2BGR)
        for c in circles1[0, :]:
            try:
                if len(c) != 3: continue
            except TypeError:
                continue
            c = list(map(lambda x: int(x), c))
            # print(c)
            cv2.circle(draw1, (c[0], c[1]), 1, (0, 255, 0), 2)
            cv2.circle(draw1, (c[0], c[1]), c[2], (0, 128, 255), 2)
        helper.Show.imshow(win_name, draw1, 2)
        return draw1

    def onchange_param3(win_name: str, image: np.ndarray, val: int):
        circles1 = cv2.HoughCircles(im, cv2.HOUGH_GRADIENT, 2, im.shape[1] / 8, 70, 80, minRadius=val, maxRadius=0)
        # print(circles)
        draw1 = cv2.cvtColor(image.copy(), cv2.COLOR_GRAY2BGR)
        for c in circles1[0, :]:
            try:
                if len(c) != 3: continue
            except TypeError:
                continue
            c = list(map(lambda x: int(x), c))
            # print(c)
            cv2.circle(draw1, (c[0], c[1]), 1, (0, 255, 0), 2)
            cv2.circle(draw1, (c[0], c[1]), c[2], (0, 128, 255), 2)
        helper.Show.imshow(win_name, draw1, 2)
        return draw1

    def onchange_param4(win_name: str, image: np.ndarray, val: int):
        circles1 = cv2.HoughCircles(im, cv2.HOUGH_GRADIENT, 2, im.shape[1] / 8, 70, 80, minRadius=30, maxRadius=val)
        # print(circles)
        draw1 = cv2.cvtColor(image.copy(), cv2.COLOR_GRAY2BGR)
        for c in circles1[0, :]:
            try:
                if len(c) != 3: continue
            except TypeError:
                continue
            c = list(map(lambda x: int(x), c))
            # print(c)
            cv2.circle(draw1, (c[0], c[1]), 1, (0, 255, 0), 2)
            cv2.circle(draw1, (c[0], c[1]), c[2], (0, 128, 255), 2)
        helper.Show.imshow(win_name, draw1, 2)
        return draw1

    # helper.Controls.adjust_multi(Configs.window_name + "圆调参", im, [
    #     helper.Controls.Arg("参数1", onchange_param1, 0, 120),
    #     helper.Controls.Arg("参数2", onchange_param2, 0, 70),
    # ])

    if use_track:
        helper.Controls.adjust_multi(Configs.window_name + "圆调参", canny, [
            helper.Controls.Arg("param1", onchange_param1, 70, 380),
            helper.Controls.Arg("param2", onchange_param2, 220, 380),
            helper.Controls.Arg("minRadius", onchange_param3, 30, 360),
            helper.Controls.Arg("maxRadius", onchange_param4, 0, 360),
        ])

    # # circles = cv2.HoughCircles(im, cv2.HOUGH_GRADIENT, 2, im.shape[1] / 8, 70, 220, minRadius=30, maxRadius=0)
    # circles = cv2.HoughCircles(im, cv2.HOUGH_GRADIENT, 2, im.shape[1] / 8, minRadius=30, maxRadius=0)
    # # print(circles)
    # if circles is None or len(circles) == 0:
    #     return src

    # 从不完美到完美进行搜索，找唯一一个
    need_search = True
    circles = circles = cv2.HoughCircles(im, cv2.HOUGH_GRADIENT, 2, im.shape[1] / 8, 70, 220, minRadius=30, maxRadius=0)
    circles_ = []
    if circles is not None:
        if len(circles) == 1 and len(circles[0]) > 1:
            circles = circles[0]
        if len(circles) == 1:
            logger.debug("Needn't search!")
            need_search = False
    else:
        need_search = True
    if need_search:
        for p2 in range(90, 380, 20):
            # print('searching', p2)
            circles = cv2.HoughCircles(im, cv2.HOUGH_GRADIENT, 2, im.shape[1] / 8, 100, p2, minRadius=30, maxRadius=0)
            if circles is None:
                break
            if len(circles) == 4 and len(circles[0]) == 1:
                continue
            if len(circles) == 1 and len(circles[0]) > 1:
                circles = circles[0]
            # print('got %s circles', len(circles))
            # if len(circles) <= 5:
            # print('circles are:', circles)
            if len(circles) == 1:
                break
            if len(circles) == 0:
                circles = circles_
                break
            circles_ = circles
        if circles.shape == (4, 1):
            for p2 in range(70, 10, -20):
                # print('searching', p2)
                circles = cv2.HoughCircles(im, cv2.HOUGH_GRADIENT, 2, im.shape[1] / 8, 100, p2, minRadius=30,
                                           maxRadius=0)
                if circles is None:
                    break
                if len(circles) == 4 and len(circles[0]) == 1:
                    continue
                if len(circles) == 1 and len(circles[0]) > 1:
                    circles = circles[0]
                # print('got %s circles', len(circles))
                # if len(circles) <= 5:
                # print('circles are:', circles)
                if len(circles) == 1:
                    break
                if len(circles) == 0:
                    circles = circles_
                    break
                circles_ = circles
    if circles is None:
        return src
    # print(circles)
    for c in circles[0, :]:
        try:
            if len(c) != 3: continue
        except TypeError:
            continue
        c = list(map(lambda x: int(x), c))
        # print(c)
        # cv2.circle(draw, (c[0], c[1]), 1, (0, 255, 0), 2)
        cv2.rectangle(draw, (c[0] - c[2], c[1] - c[2]), (c[0] + c[2], c[1] + c[2]), color_rect, 3)
        cv2.circle(draw, (c[0], c[1]), c[2], color_circle, 5)
    if len(circles) == 1:
        logger.warning("Circle at: %s" % str([x[0] for x in np.array(circles).reshape((3, 1)).tolist()]))
    helper.Show.imshow(Configs.window_name + "球", draw, 1)
    helper.Controls.wait_exit(0)
    return src


def main():
    if not os.path.exists(Configs.file_path):
        raise FileNotFoundError
    for filename in os.listdir(Configs.file_path):
        if filename.split('.')[-1].lower() not in Configs.file_exts:
            continue
        logger.info("preparing %s..." % filename)
        src = cv2.imread(os.path.join(Configs.file_path, filename))
        logger.info("got image %s, %s" % (src.shape, src.dtype))
        logger.info('processing...')
        process(src, use_track=False)


def test_speed():
    cap = cv2.VideoCapture(0)
    if not cap.isOpened():
        raise Exception("Can't open camera0!")
    while True:
        _, im = cap.read()
        resize = 4
        im = cv2.resize(im, (im.shape[1] // resize, im.shape[0] // resize))
        # 关掉按键等待
        helper.HelperConfig.wait_time_forever = 1
        process(im, refresh=False)
        helper.Show.fps_log()


if __name__ == '__main__':
    try:
        main()
        test_speed()
    except Exception:
        traceback.print_exc()
        cv2.destroyAllWindows()
        helper.Controls.exit_kill()
